# coding=utf-8
from __future__ import print_function, absolute_import
from gm.api import *
import pandas as pd
import numpy as np
import datetime
import math

# 策略中必须有init方法
def init(context):
    context.holding_num = 10# 持股数量
    context.max_drawdown = -0.05# 最大回撤
    context.holding_info = {}# 持仓股信息
    schedule(schedule_func=algo, date_rule='1d', time_rule='09:30:00')


def algo(context):
    # 上一交易日
    pre_date = get_previous_N_trading_date(context.now)
    # 获取A股代码（剔除新股、ST）
    stocks_all,stocks_all_str = get_normal_stocks(pre_date)
    
    # 通过“守株待兔”指标选股
    stocks2buy = waiting_rabbit(context,stocks_all_str,pre_date)
    print('{}:股票池：{}'.format(context.now,stocks2buy))
    
    for stock in stocks2buy:
        # 所有持仓
        Account_positions = context.account().positions()
        holding_num = len(Account_positions)
        if holding_num<context.holding_num:
            new_price = current(symbols=stock)[0]['price']
            order_target_percent(stock, percent=1/context.holding_num, position_side=1, order_type=OrderType_Limit, price=new_price)
            context.holding_info[stock] = {'high':new_price}
            subscribe(symbols=stock, frequency='60s', count=1)


def on_bar(context, bars):
    stock = bars[0]['symbol']
    new_price = bars[0]['close']
    context.holding_info[stock]['high'] = max(context.holding_info[stock]['high'],new_price)
    if new_price/context.holding_info[stock]['high']-1<=context.max_drawdown:
        order_target_percent(stock, percent=0, position_side=1, order_type=OrderType_Limit, price=new_price)
        context.holding_info.pop(stock)
        unsubscribe(symbols=stock, frequency='60s')


def waiting_rabbit(context,security,date):
    """
    守株待兔选股指标
    ：param security：股票池（格式为逗号隔开的字符串：'***,***'）
    ：param date：日期
    """
    start_date = get_previous_N_trading_date(date,counts=65+1)
    
    close = history_new(security=security,frequency='1d',start_time=start_date,end_time=date,fields='symbol,eob,close',skip_suspended=True,fill_missing=None,adjust=ADJUST_PREV,adjust_end_time=None, df=True, type=True).fillna(method='ffill')
    open_ = history_new(security=security,frequency='1d',start_time=start_date,end_time=date,fields='symbol,eob,open',skip_suspended=True,fill_missing=None,adjust=ADJUST_PREV,adjust_end_time=None, df=True, type=True).fillna(method='ffill')
    high = history_new(security=security,frequency='1d',start_time=start_date,end_time=date,fields='symbol,eob,high',skip_suspended=True,fill_missing=None,adjust=ADJUST_PREV,adjust_end_time=None, df=True, type=True).fillna(method='ffill')
    low = history_new(security=security,frequency='1d',start_time=start_date,end_time=date,fields='symbol,eob,low',skip_suspended=True,fill_missing=None,adjust=ADJUST_PREV,adjust_end_time=None, df=True, type=True).fillna(method='ffill')
    volume = history_new(security=security,frequency='1d',start_time=start_date,end_time=date,fields='symbol,eob,volume',skip_suspended=True,fill_missing=None,adjust=ADJUST_PREV,adjust_end_time=None, df=True, type=True).fillna(method='ffill')
    # 角05
    ema5 = close.apply(EMA,N=5)
    angle05 = series_atan((((ema5/ema5.shift(1)-1)*100)).iloc[-1,:])*57.3
    # 角10
    ema10 = close.apply(EMA,N=10)
    angle10 = series_atan((((ema10/ema10.shift(1)-1)*100)).iloc[-1,:])*57.3
    # 角20
    ema20 = close.apply(EMA,N=20)
    angle20 = series_atan((((ema20/ema20.shift(1)-1)*100)).iloc[-1,:])*57.3
    # A1
    A1 = close.rolling(window=5).mean()
    # A2
    A2 = close.rolling(window=10).mean()
    # A3
    A3 = close.rolling(window=20).mean()
    # A4
    A4 = close.rolling(window=60).mean()
    # CV
    CV = close.copy()
    # BBIBOLL:=(MA(CV,3)+MA(CV,6)+MA(CV,12)+MA(CV,24))/4;
    BBIBOLL = (close.rolling(window=3).mean()+close.rolling(window=6).mean()+close.rolling(window=12).mean()+close.rolling(window=24).mean())/4
    # UPR:=BBIBOLL+6*STD(BBIBOLL,11);
    UPR = BBIBOLL.iloc[-1,:]+6*BBIBOLL.rolling(window=11,min_periods=11).std()
    # QJJ:=VOL/((HIGH-LOW)*2-ABS(CLOSE-OPEN));
    QJJ = volume/((high-low)*2-abs(close-open_))
    # XVL:=(IF(CLOSE>OPEN,QJJ*(HIGH-LOW),IF(CLOSE<OPEN,QJJ*(HIGH-OPEN+CLOSE-LOW),VOL/2))+IF(CLOSE>OPEN,0-QJJ*(HIGH-CLOSE+OPEN-LOW),IF(CLOSE<OPEN,0-QJJ*(HIGH-LOW),0-VOL/2)));
    XVL = np.where(close>open_,QJJ*(high-low),np.where(close<open_,QJJ*(high-open_+close-low),volume/2))+np.where(close>open_,-QJJ*(high-close+open_-low),np.where(close<open_,-QJJ*(high-low),-volume/2))
    XVL = pd.DataFrame(XVL,index=QJJ.index,columns=QJJ.columns)
    # HSL:=(XVL/3)/1.15;
    HSL = (XVL/3)/1.15
    # 攻击:=((HSL*0.55+(REF(HSL,1)*0.33))+(REF(HSL,2)*0.22));
    attack = (HSL*0.55+(HSL.shift(1)*0.33)+HSL.shift(2)*0.22)
    # LLJX:=EMA(攻击,3);
    LLJX = attack.apply(EMA,N=3)
    # B1:=IF(SUM(LLJX,1)>REF(MA(SUM(LLJX,1),1),1),SUM(LLJX,1),DRAWNULL);
    LLJX_sum1 = LLJX.rolling(window=1,min_periods=1).sum()
    B1 = np.where(LLJX_sum1.iloc[-1,:]>LLJX_sum1.iloc[-2,:],True,False)
    # B1 = np.where(LLJX_sum1.iloc[-1,:]>LLJX_sum1.iloc[-2,:],LLJX_sum1.iloc[-1,:],False)
    B1 = pd.DataFrame(B1,index=LLJX.columns)
    B1 = B1[B1.columns[0]]
    # B2:=IF(SUM(LLJX,2)>REF(MA(SUM(LLJX,2),1),1),SUM(LLJX,2),DRAWNULL);
    LLJX_sum2 = LLJX.rolling(window=2,min_periods=2).sum()
    B2 = np.where(LLJX_sum2.iloc[-1,:]>LLJX_sum2.iloc[-2,:],True,False)
    # B2 = np.where(LLJX_sum2.iloc[-1,:]>LLJX_sum2.iloc[-2,:],LLJX_sum2.iloc[-1,:],False)
    B2 = pd.DataFrame(B2,index=LLJX.columns)
    B2 = B2[B2.columns[0]]
    # B3:=IF(SUM(LLJX,3)>REF(MA(SUM(LLJX,3),1),1),SUM(LLJX,3),DRAWNULL);
    LLJX_sum3 = LLJX.rolling(window=3,min_periods=3).sum()
    B3 = np.where(LLJX_sum3.iloc[-1,:]>LLJX_sum3.iloc[-2,:],True,False)
    # B3 = np.where(LLJX_sum3.iloc[-1,:]>LLJX_sum3.iloc[-2,:],LLJX_sum3.iloc[-1,:],False)
    B3 = pd.DataFrame(B3,index=LLJX.columns)
    B3 = B3[B3.columns[0]]
    # B4:=IF(SUM(LLJX,4)>REF(MA(SUM(LLJX,4),1),1),SUM(LLJX,4),DRAWNULL);
    LLJX_sum4 = LLJX.rolling(window=4,min_periods=4).sum()
    B4 = np.where(LLJX_sum4.iloc[-1,:]>LLJX_sum4.iloc[-2,:],True,False)
    # B4 = np.where(LLJX_sum4.iloc[-1,:]>LLJX_sum4.iloc[-2,:],LLJX_sum4.iloc[-1,:],False)
    B4 = pd.DataFrame(B4,index=LLJX.columns)
    B4 = B4[B4.columns[0]]
    # B5:=IF(SUM(LLJX,5)>REF(MA(SUM(LLJX,5),1),1),SUM(LLJX,5),DRAWNULL);
    LLJX_sum5 = LLJX.rolling(window=5,min_periods=5).sum()
    B5 = np.where(LLJX_sum5.iloc[-1,:]>LLJX_sum5.iloc[-2,:],True,False)
    # B5 = np.where(LLJX_sum5.iloc[-1,:]>LLJX_sum5.iloc[-2,:],LLJX_sum5.iloc[-1,:],False)
    B5 = pd.DataFrame(B5,index=LLJX.columns)
    B5 = B5[B5.columns[0]]
    # B6:=IF(SUM(LLJX,6)>REF(MA(SUM(LLJX,6),1),1),SUM(LLJX,6),DRAWNULL);
    LLJX_sum6 = LLJX.rolling(window=6,min_periods=6).sum()
    B6 = np.where(LLJX_sum6.iloc[-1,:]>LLJX_sum6.iloc[-2,:],True,False)
    # B6 = np.where(LLJX_sum6.iloc[-1,:]>LLJX_sum6.iloc[-2,:],LLJX_sum6.iloc[-1,:],False)
    B6 = pd.DataFrame(B6,index=LLJX.columns)
    B6 = B6[B6.columns[0]]
    # B7:=IF(SUM(LLJX,7)>REF(MA(SUM(LLJX,7),1),1),SUM(LLJX,7),DRAWNULL);
    LLJX_sum7 = LLJX.rolling(window=7,min_periods=7).sum()
    B7 = np.where(LLJX_sum7.iloc[-1,:]>LLJX_sum7.iloc[-2,:],True,False)
    # B7 = np.where(LLJX_sum7.iloc[-1,:]>LLJX_sum7.iloc[-2,:],LLJX_sum7.iloc[-1,:],False)
    B7 = pd.DataFrame(B7,index=LLJX.columns)
    B7 = B7[B7.columns[0]]
    # B8:=IF(SUM(LLJX,8)>REF(MA(SUM(LLJX,8),1),1),SUM(LLJX,8),DRAWNULL);
    LLJX_sum8 = LLJX.rolling(window=8,min_periods=8).sum()
    B8 = np.where(LLJX_sum8.iloc[-1,:]>LLJX_sum8.iloc[-2,:],True,False)
    # B8 = np.where(LLJX_sum8.iloc[-1,:]>LLJX_sum8.iloc[-2,:],LLJX_sum8.iloc[-1,:],False)
    B8 = pd.DataFrame(B8,index=LLJX.columns)
    B8 = B8[B8.columns[0]]
    # B9:=IF(SUM(LLJX,9)>REF(MA(SUM(LLJX,9),1),1),SUM(LLJX,9),DRAWNULL);
    LLJX_sum9 = LLJX.rolling(window=9,min_periods=9).sum()
    B9 = np.where(LLJX_sum9.iloc[-1,:]>LLJX_sum9.iloc[-2,:],True,False)
    # B9 = np.where(LLJX_sum9.iloc[-1,:]>LLJX_sum9.iloc[-2,:],LLJX_sum9.iloc[-1,:],False)
    B9 = pd.DataFrame(B9,index=LLJX.columns)
    B9 = B9[B9.columns[0]]
    # B10:=IF(SUM(LLJX,10)>REF(MA(SUM(LLJX,10),1),1),SUM(LLJX,10),DRAWNULL);
    LLJX_sum10 = LLJX.rolling(window=10,min_periods=10).sum()
    B10 = np.where(LLJX_sum10.iloc[-1,:]>LLJX_sum10.iloc[-2,:],True,False)
    # B10 = np.where(LLJX_sum10.iloc[-1,:]>LLJX_sum10.iloc[-2,:],LLJX_sum10.iloc[-1,:],False)
    B10 = pd.DataFrame(B10,index=LLJX.columns)
    B10 = B10[B10.columns[0]]
    # XG: A1>A2 AND A2>A3 AND A1>REF(A1,1) AND A2>REF(A2,1) AND A3>REF(A3,1) 
    # AND CROSS(MA(C,5),MA(C,60)) AND C>MA(C,60) AND 角05>65 AND 角10>60 AND 角20>45 COUNT(MA(C,60)>REF(MA(C,60),1),5)>3 
    # AND B1 AND B2 AND B3 AND B4 AND B5 AND B6 AND B7 AND B8 AND B9 AND B10 ;
    stand1 = ((A1>A2).iloc[-1,:]) & ((A2>A3).iloc[-1,:]) & ((A1>A1.shift(1)).iloc[-1,:]) & ((A2>A2.shift(1)).iloc[-1,:]) & ((A3>A3.shift(1)).iloc[-1,:])
    stand2 = ((A1<A4).iloc[-2,:]) & ((A1>A4).iloc[-1,:]) & ((close>A4).iloc[-1,:]) & (angle05>65) & (angle10>60) & (angle20>45) & ((A4>A4.shift(1)).iloc[-5:,:].sum()>3)
    stand3 = (B1) & (B2) & (B3) & (B4) & (B5) & (B6) & (B7) & (B8) & (B9) & (B10)
    stand = (stand1) & (stand2) & (stand3)
    to_buy = list(stand[stand==True].index)
    return to_buy
    

def series_atan(series):
    series_new = series.copy()
    for x in series.index:
        series_new.loc[x] = math.atan(series.loc[x])
    return series_new

def get_previous_N_trading_date(date,counts=1,exchanges='SHSE'):
    """
    获取end_date前N个交易日,end_date为datetime格式，不包括date日期
    :param date：目标日期
    :param counts：历史回溯天数，默认为1，即前一天
    """
    date = pd.Timestamp(date)
    previous_N_trading_date = get_trading_dates(exchange=exchanges, 
                                                start_date=date-datetime.timedelta(days=max(counts+30,counts*3)), 
                                                end_date=date)[-counts-1]
    return previous_N_trading_date


def get_normal_stocks(date,new_days=365):
    """
    获取目标日期date的A股代码（剔除停牌股、ST股、次新股（365天））
    :param date：目标日期
    :param new_days:新股上市天数，默认为365天
    """
    if isinstance(date,str) and len(date)==10:
        date = datetime.datetime.strptime(date,"%Y-%m-%d")
    elif isinstance(date,str) and len(date)>10:
        date = datetime.datetime.strptime(date,"%Y-%m-%d %H:%M:%S")
    date = date.date()
    # 先剔除退市股、次新股和B股
    df_code = get_instrumentinfos(sec_types=SEC_TYPE_STOCK, fields='symbol, listed_date, delisted_date', df=True)
    df_code['listed_date'] = df_code['listed_date'].apply(lambda x:x.date())
    df_code['delisted_date'] = df_code['delisted_date'].apply(lambda x:x.date())
    all_stocks = [code for code in df_code[(df_code['listed_date']<=date-datetime.timedelta(days=new_days))&(df_code['delisted_date']>date)].symbol.to_list() if code[:6]!='SHSE.9' and code[:6]!='SZSE.2']
    # 再剔除当前的停牌股和ST股
    history_ins = get_history_instruments(symbols=all_stocks, start_date=date, end_date=date, fields='symbol,sec_level, is_suspended', df=True)
    all_stocks = list(history_ins[(history_ins['sec_level']==1) & (history_ins['is_suspended']==0)]['symbol'])
    all_stocks_str = ','.join(all_stocks)
    return all_stocks,all_stocks_str


def history_new(security,frequency,start_time,end_time,fields,skip_suspended=True,fill_missing=None,adjust=ADJUST_PREV,adjust_end_time=None, df=True, type=True):
    """
    分区间获取数据（以避免超出数据限制）(start_time和end_date为字符串,fields需包含eob和symbol,单字段)
    :param ：参数同history()参数一致，adjust_end_time默认为回测结束时间：None,注意需要根据不同场景使用end_time或context.backtest_end_time
    :param type：默认为True，输出2维DataFrame（日期*股票）,否则输出1维DataFrame
    """
    Data = pd.DataFrame()
    if frequency=='1d':
        trading_date = pd.Series(get_trading_dates(exchange='SZSE', start_date=start_time, end_date=end_time))
    else:
        trading_date = history('SHSE.000300', frequency=frequency, start_time=start_time, end_time=end_time, fields='eob', skip_suspended=skip_suspended, fill_missing=fill_missing, adjust=adjust, adjust_end_time=adjust_end_time, df=df)
        trading_date = trading_date['eob']
    space = 5
    if len(trading_date)<=space:
        Data = history(security, frequency=frequency, start_time=start_time, end_time=end_time, fields=fields, skip_suspended=skip_suspended, fill_missing=fill_missing, adjust=adjust, adjust_end_time=adjust_end_time, df=df)
    else:
        for n in range(int(np.ceil(len(trading_date)/space))):
            start = n*space
            end = start+space-1
            if end>=len(trading_date):
                data = history(security, frequency=frequency, start_time=trading_date.iloc[start], end_time=trading_date.iloc[-1], fields=fields, skip_suspended=skip_suspended, fill_missing=fill_missing, adjust=adjust, adjust_end_time=adjust_end_time, df=df)
            else:
                data = history(security, frequency=frequency, start_time=trading_date.iloc[start], end_time=trading_date.iloc[end], fields=fields, skip_suspended=skip_suspended, fill_missing=fill_missing, adjust=adjust, adjust_end_time=adjust_end_time, df=df)
            if len(data)>=33000:
                print('请检查返回数据量，可能超过系统限制，缺少数据！！！！！！！！！！')
            Data = pd.concat([Data,data]) 
    Data.drop_duplicates(subset=['eob','symbol'],keep='first',inplace=True)
    if type:
        if len(Data)>0:
            Data = Data.set_index(['eob','symbol'])
            Data = Data.unstack()
            Data.columns = Data.columns.droplevel(level=0)
    return Data


def EMA(series, N=10):
    """
    Y=EMA(X,N)=[2X+(N-1)*Y']/(N+1),Y'为Y的前值
    """
    data = []
    price = np.nan_to_num(series)
    for i in range(len(price)):
        if i==0:
            data+=[price[i]]
        else:
            data+=[(2*price[i]+(N-1)*data[-1])/(N+1)]
    return data


if __name__ == '__main__':
    '''
        strategy_id策略ID, 由系统生成
        filename文件名, 请与本文件名保持一致
        mode运行模式, 实时模式:MODE_LIVE回测模式:MODE_BACKTEST
        token绑定计算机的ID, 可在系统设置-密钥管理中生成
        backtest_start_time回测开始时间
        backtest_end_time回测结束时间
        backtest_adjust股票复权方式, 不复权:ADJUST_NONE前复权:ADJUST_PREV后复权:ADJUST_POST
        backtest_initial_cash回测初始资金
        backtest_commission_ratio回测佣金比例
        backtest_slippage_ratio回测滑点比例
        '''
    run(strategy_id='8c0ef19f-9307-11ed-8edf-f46b8c02346f',
        filename='main.py',
        mode=MODE_BACKTEST,
        token='ca9c4532786a122aa5ceb3fd726432862e046f91',
        backtest_start_time='2023-01-01 08:00:00',
        backtest_end_time='2023-01-18 16:00:00',
        backtest_adjust=ADJUST_PREV,
        backtest_initial_cash=10000000,
        backtest_commission_ratio=0.0007,
        backtest_slippage_ratio=0.00123)

